% calculate_total_time.m
% 根据文档附录C的规则，计算任务的总执行时间。
% (版本 28: 增加了 code_map 作为输入参数)

function total_time = calculate_total_time(solution, graph, code_map)
    
    fprintf('--- 开始计算总执行时间 ---\n');

    % --- 准备工作: 预处理节点信息 ---
    
    % 创建一个Map，用于快速查找每个BufId对应的ALLOC和FREE节点ID
    buf_id_to_nodes = containers.Map('KeyType', 'int32', 'ValueType', 'any');
    for i = 1:length(graph.nodes)
        node = graph.nodes(i);
        if isempty(node.id) || isempty(node.buf_id), continue; end
        
        if ~isKey(buf_id_to_nodes, node.buf_id)
            buf_id_to_nodes(node.buf_id) = struct('alloc_node', [], 'free_node', []);
        end
        
        temp = buf_id_to_nodes(node.buf_id);
        if node.op_code == code_map.ALLOC
            temp.alloc_node = node.id;
        elseif node.op_code == code_map.FREE
            temp.free_node = node.id;
        end
        buf_id_to_nodes(node.buf_id) = temp;
    end

    % --- 步骤 2: 计算并添加缓存复用依赖 ---
    
    fprintf('步骤 2/3: 计算缓存复用依赖...\n');
    % 复制一份前驱列表，我们将在此基础上添加新的依赖边
    predecessors_list = graph.predecessors_list;
    
    % 模拟内存分配，以找出地址复用关系
    active_allocations = containers.Map('KeyType','int32','ValueType','any'); % key: buf_id, val: {offset, size, free_node_id}
    address_history = containers.Map('KeyType','double','ValueType','any'); % key: address, val: last_free_node_id
    
    for step = 1:length(solution.final_schedule)
        node_id = solution.final_schedule(step);
        node = graph.nodes(node_id);

        % 处理分配事件
        if node.op_code == code_map.ALLOC || strcmp(node.op, 'SPILL_IN')
            buf_id = node.buf_id;
            if isempty(buf_id), continue; end

            if node.op_code == code_map.ALLOC
                offset = solution.memory_offsets(buf_id);
            else % SPILL_IN
                offset = solution.spill_map(node_id);
            end

            % 查找此地址空间之前被哪个FREE节点释放
            last_free_node = [];
            if isKey(address_history, offset)
                last_free_node = address_history(offset);
            end

            % 如果找到了，并且这个依赖不存在，则添加
            if ~isempty(last_free_node) && ~ismember(last_free_node, predecessors_list{node.id})
                predecessors_list{node.id} = [predecessors_list{node.id}, last_free_node];
            end
            
            % 记录当前分配
            if isKey(buf_id_to_nodes, buf_id)
                node_info = buf_id_to_nodes(buf_id);
                active_allocations(buf_id) = {offset, node.size, node_info.free_node};
            end
        
        % 处理释放事件
        elseif node.op_code == code_map.FREE || strcmp(node.op, 'SPILL_OUT')
            buf_id = node.buf_id;
            if isempty(buf_id) || ~isKey(active_allocations, buf_id), continue; end

            alloc_info = active_allocations(buf_id);
            offset = alloc_info{1};
            
            % 更新地址历史，记录这个地址最后是由当前节点释放的
            address_history(offset) = node.id;
            remove(active_allocations, buf_id);
        end
    end
    
    % --- 步骤 3: 流水线排布 ---
    
    fprintf('步骤 3/3: 模拟硬件流水线排布...\n');
    
    node_end_times = zeros(graph.num_nodes, 1);
    pipe_free_time = containers.Map('KeyType', 'char', 'ValueType', 'double');
    
    % 初始化所有硬件单元的空闲时间为0
    for i = 1:graph.num_nodes
        node = graph.nodes(i);
        if ~isempty(node.pipe) && ~isKey(pipe_free_time, node.pipe)
            pipe_free_time(node.pipe) = 0;
        end
    end

    % 严格按照最终调度顺序进行模拟
    for i = 1:length(solution.final_schedule)
        node_id = solution.final_schedule(i);
        node = graph.nodes(node_id);
        
        % 1. 计算来自前序节点的依赖时间
        predecessors = predecessors_list{node_id};
        predecessor_finish_time = 0;
        if ~isempty(predecessors)
            predecessor_finish_time = max(node_end_times(predecessors));
        end
        
        % 缓存管理节点不耗时，但需要遵守依赖
        if isempty(node.cycles) || node.cycles == 0
            node_end_times(node_id) = predecessor_finish_time;
            continue;
        end
        
        % 2. 计算硬件资源可用时间
        pipe = node.pipe;
        resource_available_time = pipe_free_time(pipe);
        
        % 3. 确定节点的实际开始时间
        start_time = max(predecessor_finish_time, resource_available_time);
        
        % 4. 计算结束时间并更新状态
        end_time = start_time + node.cycles;
        node_end_times(node_id) = end_time;
        pipe_free_time(pipe) = end_time;
    end
    
    % 总执行时间是所有节点结束时间的最大值
    total_time = max(node_end_times);
    fprintf('总执行时间计算完成。\n');
end